import torch
from torch import nn
from d2l import torch as d2l

def comp_conv2d(conv2d,X):
    X = X.reshape((1,1)+X.shape)
    Y = conv2d(X)
    return Y.reshape(Y.shape[2:])

conv2d = nn.Conv2d(1,1,kernel_size=3,padding=1,stride=2)
X = torch.rand(size=(8,8))
print(comp_conv2d(conv2d,X).shape)

pool2d = nn.MaxPool2d(3)
XX = torch.arange(16,dtype=torch.float32).reshape((1,1,4,4))
print(pool2d(XX))

    